package com.flyxiaozhu.erp.admin.authorize;

import com.flyxiaozhu.erp.core.libs.ClassScanner;
import com.flyxiaozhu.erp.core.utils.CollectionUtils;
import org.apache.commons.lang.StringUtils;
import org.springframework.security.access.annotation.Secured;
import org.springframework.security.access.prepost.PostAuthorize;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.RestController;

import java.lang.reflect.Method;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class PermissionScanner extends ClassScanner {
    private String basePackage;

    private Set<Class<?>> controllers;

    public PermissionScanner(String basePackage) {
        this.basePackage = basePackage;
    }

    public Set<Class<?>> getControllers() {
        if (controllers == null) {
            controllers = new HashSet<>();
            Set<Class<?>> clsList = getClasses(basePackage, true);
            if (clsList != null && clsList.size() > 0) {
                for (Class<?> cls : clsList) {
                    if (cls.getAnnotation(Controller.class) != null || cls.getAnnotation(RestController.class) != null) {
                        controllers.add(cls);
                    }
                }
            }
        }
        return controllers;
    }

    public Set<String> getPermissions() {
        Set<String> permissions = new HashSet<>();
        for (Class<?> cls : getControllers()) {
            Method[] methods = cls.getMethods();
            for (Method method : methods) {
                Secured securedAnnotation = method.getAnnotation(Secured.class);
                if (securedAnnotation != null) {
                    String[] value = securedAnnotation.value();
                    permissions.addAll(CollectionUtils.map(value, o -> o.replaceAll("^ROLE_", "")));
                } else {
                    String expression = "";
                    PreAuthorize preAuthorizeAnnotation = method.getAnnotation(PreAuthorize.class);
                    if (preAuthorizeAnnotation != null) {
                        expression = preAuthorizeAnnotation.value();
                    }

                    PostAuthorize postAuthorizeAnnotation = method.getAnnotation(PostAuthorize.class);
                    if (postAuthorizeAnnotation != null) {
                        expression = postAuthorizeAnnotation.value();
                    }

                    if (!StringUtils.isBlank(expression)) {
                        Pattern pattern = Pattern.compile("ROLE_\\w+");
                        Matcher matcher = pattern.matcher(expression);
                        while (matcher.find()) {
                            permissions.add(matcher.group().replaceAll("^ROLE_", ""));
                        }
                    }
                }
            }
        }

        return permissions;
    }
}
